// THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY OF
// ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A
// PARTICULAR PURPOSE.
//
// Copyright (c) Microsoft Corporation. All rights reserved

#include "Private.h"
#include "RegKey.h"

//---------------------------------------------------------------------
//
// ctor
//
//---------------------------------------------------------------------

CRegKey::CRegKey() {
	_keyHandle = nullptr;
}

//---------------------------------------------------------------------
//
// dtor
//
//---------------------------------------------------------------------

CRegKey::~CRegKey() {
	Close();
}


//---------------------------------------------------------------------
//
// operator
//
//---------------------------------------------------------------------

HKEY CRegKey::GetHKEY() {
	return _keyHandle;
}

//---------------------------------------------------------------------
//
// Create
//
//---------------------------------------------------------------------

LONG CRegKey::Create(_In_ HKEY hKeyPresent, _In_ LPCWSTR pwszKeyName, _In_reads_opt_(255) LPWSTR pwszClass, DWORD dwOptions, REGSAM samDesired, _Inout_ LPSECURITY_ATTRIBUTES lpSecAttr, _Out_opt_ LPDWORD lpdwDisposition) {
	DWORD disposition = 0;
	HKEY keyHandle = nullptr;

	LONG res = RegCreateKeyEx(hKeyPresent, pwszKeyName, 0,
		pwszClass, dwOptions, samDesired, lpSecAttr, &keyHandle, &disposition);

	if (lpdwDisposition != nullptr) {
		*lpdwDisposition = disposition;
	}

	if (res == ERROR_SUCCESS) {
		Close();
		_keyHandle = keyHandle;
	}

	return res;
}

//---------------------------------------------------------------------
//
// Open
//
//---------------------------------------------------------------------

LONG CRegKey::Open(_In_ HKEY hKeyParent, _In_ LPCWSTR pwszKeyName, REGSAM samDesired) {
	HKEY keyHandle = nullptr;

	LONG res = RegOpenKeyEx(hKeyParent, pwszKeyName, 0, samDesired, &keyHandle);
	if (res == ERROR_SUCCESS) {
		Close();
		_keyHandle = keyHandle;
	}
	return res;
}

//---------------------------------------------------------------------
//
// Close
//
//---------------------------------------------------------------------

LONG CRegKey::Close() {
	LONG res = ERROR_SUCCESS;
	if (_keyHandle) {
		res = RegCloseKey(_keyHandle);
		_keyHandle = nullptr;
	}
	return res;
}

//---------------------------------------------------------------------
//
// DeleteSubKey
// RecurseDeleteKey
//
//---------------------------------------------------------------------

LONG CRegKey::DeleteSubKey(_In_ LPCWSTR pwszSubKey) {
	return RegDeleteKey(_keyHandle, pwszSubKey);
}

LONG CRegKey::RecurseDeleteKey(_In_ LPCWSTR pwszSubKey) {
	CRegKey key;
	LONG res = key.Open(_keyHandle, pwszSubKey, KEY_READ | KEY_WRITE);
	if (res != ERROR_SUCCESS) {
		return res;
	}
	FILETIME time;
	WCHAR subKeyName[256] = { '\0' };
	DWORD subKeyNameSize = ARRAYSIZE(subKeyName);

	while (RegEnumKeyEx(key.GetHKEY(), 0, subKeyName, &subKeyNameSize, NULL, NULL, NULL, &time) == ERROR_SUCCESS) {
		subKeyName[ARRAYSIZE(subKeyName) - 1] = L'\0';
		res = key.RecurseDeleteKey(subKeyName);
		if (res != ERROR_SUCCESS) {
			return res;
		}
		subKeyNameSize = ARRAYSIZE(subKeyName);
	}

	key.Close();

	return DeleteSubKey(pwszSubKey);
}

//---------------------------------------------------------------------
//
// DeleteValue
//
//---------------------------------------------------------------------

LONG CRegKey::DeleteValue(_In_ LPCWSTR pwszValue) {
	return RegDeleteValue(_keyHandle, pwszValue);
}

//---------------------------------------------------------------------
//
// QueryStingValue
// SetStringValue
//
//---------------------------------------------------------------------

LONG CRegKey::QueryStringValue(_In_opt_ LPCWSTR pwszValueName, _Out_writes_opt_(*pnChars) LPWSTR pwszValue, _Inout_ ULONG* pnChars) {
	LONG res = 0;
	DWORD dataType = REG_NONE;
	ULONG pwszValueSize = 0;

	if (pnChars == nullptr) {
		return E_INVALIDARG;
	}

	pwszValueSize = (*pnChars) * sizeof(WCHAR);
	*pnChars = 0;

	res = RegQueryValueEx(_keyHandle, pwszValueName, NULL, &dataType, (LPBYTE)pwszValue, &pwszValueSize);
	if (res != ERROR_SUCCESS) {
		return res;
	}
	if ((dataType != REG_SZ) && (dataType != REG_EXPAND_SZ)) {
		return ERROR_INVALID_DATA;
	}

	*pnChars = pwszValueSize / sizeof(WCHAR);

	return ERROR_SUCCESS;
}

LONG CRegKey::SetStringValue(_In_opt_ LPCWSTR pwszValueName, _In_ LPCWSTR pwszValue, DWORD dwType) {
	size_t lenOfValue = 0;
	if (pwszValue == nullptr) {
		return ERROR_INVALID_PARAMETER;
	}

	if (StringCchLength(pwszValue, STRSAFE_MAX_CCH, &lenOfValue) != S_OK) {
		return ERROR_INVALID_PARAMETER;
	}
	DWORD len = static_cast<DWORD>(lenOfValue);
	return RegSetValueEx(_keyHandle, pwszValueName, NULL, dwType, (LPBYTE)pwszValue, (++len) * sizeof(WCHAR));
}

//---------------------------------------------------------------------
//
// QueryDWORDValue
// SetDWORDValue
//
//---------------------------------------------------------------------

LONG CRegKey::QueryDWORDValue(_In_opt_ LPCWSTR pwszValueName, _Out_ DWORD& dwValue) {
	LONG res = 0;
	DWORD dataType = REG_NONE;
	ULONG pwszValueSize = 0;

	pwszValueSize = sizeof(DWORD);
	res = RegQueryValueEx(_keyHandle, pwszValueName, NULL, &dataType, (LPBYTE)(&dwValue), &pwszValueSize);
	if (res != ERROR_SUCCESS) {
		return res;
	}
	if (dataType != REG_DWORD) {
		return ERROR_INVALID_DATA;
	}

	return ERROR_SUCCESS;
}

LONG CRegKey::SetDWORDValue(_In_opt_ LPCWSTR pwszValueName, DWORD dwValue) {
	return RegSetValueEx(_keyHandle, pwszValueName, NULL, REG_DWORD, (LPBYTE)(&dwValue), sizeof(DWORD));
}

//---------------------------------------------------------------------
//
// QueryBinaryValue
// SetBinaryValue
//
//---------------------------------------------------------------------

LONG CRegKey::QueryBinaryValue(_In_opt_ LPCWSTR pwszValueName, _Out_writes_opt_(cbData) BYTE* lpData, DWORD cbData) {
	LONG res = 0;
	DWORD dataType = REG_NONE;
	ULONG pwszValueSize = 0;

	pwszValueSize = cbData;
	res = RegQueryValueEx(_keyHandle, pwszValueName, NULL, &dataType, lpData, &pwszValueSize);
	if (res != ERROR_SUCCESS) {
		return res;
	}
	if (dataType != REG_BINARY) {
		return ERROR_INVALID_DATA;
	}
	if (pwszValueSize != cbData) {
		return ERROR_INVALID_DATA;
	}

	return ERROR_SUCCESS;
}

LONG CRegKey::SetBinaryValue(_In_opt_ LPCWSTR pwszValueName, _In_reads_(cbData) BYTE* lpData, DWORD cbData) {
	return RegSetValueEx(_keyHandle, pwszValueName, NULL, REG_BINARY, lpData, cbData);
}
